import os
import cv2
import argparse

def split_augment_image(path, file, output):
    image = cv2.imread(path)
    heigh = image.shape[0]
    width = image.shape[1]
    filename = file.split(".")
    count = 0
    for w in range(0, width-256, 256):
        for h in range(0, heigh-256, 256):
            count += 1
            result_image = image[h:h+512,w:w+512]

            # [0, 1, 2, 3] * 90 , No rotation, 90 counterclockwise, 180 counterclockwise, 270 counterclockwise
            for i in range(4):
                change_image = cv2.warpAffine(result_image, cv2.getRotationMatrix2D((256,256), i*90, 1.0), (512, 512))
                cv2.imwrite("{}/{}_{}_{}.{}".format(output, filename[0], str(count).zfill(2), str(i+1), filename[1]), change_image)
            
            change_image = cv2.flip(result_image, 0)# 0 flipped vertically, 1 flipped horizontally
            cv2.imwrite("{}/{}_{}_{}.{}".format(output, filename[0], str(count).zfill(2), "5", filename[1]), change_image)
            
            change_image = cv2.flip(result_image, 1)# 0 flipped vertically, 1 flipped horizontally
            cv2.imwrite("{}/{}_{}_{}.{}".format(output, filename[0], str(count).zfill(2), "6", filename[1]), change_image)

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Split image')
    parser.add_argument('--input', '-i', type=str, required=True, help='Path to the input image')
    parser.add_argument('--output', '-o', type=str, required=True, help='Path to the output image')
    args = parser.parse_args()

    for root, dirs, files in os.walk(args.input):
        for file in files:
            if 'tif' in file:
                split_augment_image("{}/{}".format(root,file), file, args.output)